import numpy as np
import gradio as gr
from PIL import Image
from pinecone import Pinecone

# 连接到Pinecone并获取索引
def get_pinecone_index():
    """连接到Pinecone索引"""
    pc = Pinecone(api_key="pcsk_4js8wK_9HDCxEApcaiaSiAibHfiHkgeP98qpaV27QoyeqJn2tv2biGZKhN8pX9MXuxpE2H")
    index_name = "mnist-index"
    
    # 检查索引是否存在
    if index_name not in pc.list_indexes().names():
        raise ValueError(f"Pinecone索引 '{index_name}' 不存在，请先运行pinecone_train.py创建并填充索引")
    
    return pc.Index(index_name)

# 获取Pinecone索引实例
index = get_pinecone_index()

def preprocess_and_predict(image, k=11):
    """预处理    预处理输入图像并使用Pinecone进行KNN预测
    k: 近邻数量
    """
    if image is None:
        return "请先绘制一个数字"
    
    # 确保图像是PIL格式
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # 图像预处理：转为灰度图、调整尺寸、归一化
    gray_image = image.convert("L")  # 转为灰度图
    resized_image = gray_image.resize((8, 8), Image.Resampling.LANCZOS)  # 调整为8x8
    image_array = np.array(resized_image)
    
    # 反转颜色（匹配数据集的黑白对比度）并归一化到0-16范围
    inverted_array = 255 - image_array
    normalized_array = (inverted_array / 255) * 16
    
    # 调整数组形状为Pinecone查询格式
    query_vector = normalized_array.flatten().tolist()
    
    # 使用Pinecone进行KNN搜索
    results = index.query(
        vector=query_vector,
        top_k=k,
        include_metadata=True
    )
    
    # 提取邻居标签并投票决定预测结果
    labels = [match['metadata']['label'] for match in results['matches']]
    most_common = max(set(labels), key=labels.count)
    
    return f"预测结果：{most_common}"

# 创建Gradio界面
interface = gr.Interface(
    fn=preprocess_and_predict,
    inputs=[
        gr.Sketchpad(
            height=400, 
            width=400, 
            label="请在此绘制数字",
            brush_radius=2  # 画笔粗细
        ),
        gr.Slider(minimum=1, maximum=31, value=11, step=2, 
                 label="K值（近邻数量）")
    ],
    outputs="text",
    title="基于Pinecone的KNN手写数字识别",
    description="请在画板上绘制0-9的数字，系统将使用Pinecone云服务进行识别"
)

if __name__ == "__main__":
    interface.launch(share=False)